import subprocess
import time
import json


def run_stress_test(cpu_load, vm_bytes, stress_run_time):
    text = f'---启动压力测试，cpu压力{cpu_load}%，内存压力{vm_bytes}%，运行{stress_run_time}s---'
    cmd = [
        'kubectl',
        'exec',
        'stress-ng-deployment-0620-67bd678847-nhblv',
        '--',
        'stress-ng',
        '--cpu', '16',
        '--cpu-method', 'pi',
        '--cpu-load', str(cpu_load),
        '--vm', '1',
        '--vm-bytes', str(vm_bytes) + '%',
        '--vm-keep', '-t', str(stress_run_time) + 's'
    ]
    proc = subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    time.sleep(45)
    return proc


def run_stress_test_v2(cpu_load, stress_run_time):
    text = f'---启动压力测试，cpu压力{cpu_load}%，运行{stress_run_time}s---'
    cmd = [
        'kubectl',
        'exec',
        'stress-ng-deployment-0620-67bd678847-nhblv',
        '--',
        'stress-ng',
        '--cpu', '16',
        '--cpu-method', 'pi',
        '--cpu-load', str(cpu_load),
        '-t', str(stress_run_time) + 's'
    ]
    proc = subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    time.sleep(45)
    return proc


def node_status_check():
    cmd = ['kubectl', 'top', 'nodes']
    subprocess.run(cmd, check=True)


def run_redis_bench_test(request_num, client_num):
    overrides = {
        "spec": {
            "nodeName": "fuyao-master"
        }
    }
    overrides_json = json.dumps(overrides)
    
    cmd = [
        'kubectl',
        'run',
        'redis-test-0623',
        '--image=redis:latest',
        f'--overrides={overrides_json}',
        '--restart=Never',
        '-it',
        '--',
        'redis-benchmark',
        '-h', 'redis-service',
        '-p', '6379',
        '-n', str(request_num),
        '-c', str(client_num),
        '-t', 'get'
    ]
    subprocess.run(cmd, check=True)


def clean_up(stress_proc):
    cmd_delete_redis_container = [
        'kubectl',
        'delete',
        'pod',
        'redis-test-0623'
    ]
    subprocess.run(cmd_delete_redis_container, check=False)
    if stress_proc:
        stress_proc.terminate()
        time.sleep(1)
        if stress_proc.poll() is None:
            stress_proc.kill()


def rps_change_pipeline():
    # 开始业务服务QPS变化测试
    # 输入1为无压环境，输入2为有压环境：
    stress_proc = ''
    test_mode_list = ['No Stress', 'With Stress']
    mode_input = input()
    test_mode = test_mode_list[int(mode_input) - 1]
    if test_mode in 'With Stress':
        stress_proc = run_stress_test(100, 90, 360)
    node_status_check()
    run_redis_bench_test(10000000, 100)
    node_status_check()
    clean_up(stress_proc)


def rps_cpu_pipeline():
    # 开始节点CPU利用率提升测试...
    # 0为无压环境，1-10分别为cpu压力10%~70%环境：
    stress_proc = ''
    test_mode_list = ['No Stress', 'cpu10%', 'cpu20%', 'cpu30%', 'cpu40%', 'cpu50%', 'cpu60%', 'cpu70%']
    mode_input = input()
    cpu_load = int(mode_input) * 10
    test_mode = test_mode_list[int(mode_input)]
    if test_mode not in 'No Stress':
        stress_proc = run_stress_test_v2(cpu_load, 360)
    node_status_check()
    run_redis_bench_test(10000000, 100)
    node_status_check()
    clean_up(stress_proc)


def rps_mem_pipeline():
    # 开始节点内存利用率提升测试...
    # 0为无压环境，1-5分别为mem压力25%~75%环境：
    ref_dict = {}
    ref_dict[1] = 25
    ref_dict[2] = 37.5
    ref_dict[3] = 50
    ref_dict[4] = 62.5
    ref_dict[5] = 75
    stress_proc = ''
    test_mode_list = ['No Stress', 'mem25%', 'mem37.5%', 'mem50%', 'mem62.5%', 'mem75%']
    mode_input = input()
    vm_bytes = ref_dict.get(int(mode_input), 0)
    test_mode = test_mode_list[int(mode_input)]
    if test_mode not in 'No Stress':
        stress_proc = run_stress_test(40, vm_bytes, 360)
    node_status_check()
    run_redis_bench_test(10000000, 100)
    node_status_check()
    clean_up(stress_proc)



if __name__ == '__main__':

    # 业务服务QPS变化
    rps_change_pipeline()

    # 节点CPU利用率提升
    # rps_cpu_pipeline()

    # 节点内存利用率提升
    # rps_mem_pipeline()
